import numpy as np
from cv2 import cv2
import re
import matplotlib.pyplot as plt
from PIL import Image
from torchvision import transforms
import torch
import torchvision.datasets as datasets
from utils2 import *
import random
from sklearn.cluster import KMeans
from sklearn import decomposition
import scipy.stats
import matplotlib.pyplot as plt
from scipy.stats import wasserstein_distance
from sklearn.metrics import confusion_matrix
import pickle


video_path1 = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round1/round1.avi'
video_path2 = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/round2.avi'
video_path3 = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round3/round3.avi'
index_map = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/IndexMap.pickle'
position_path1 = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round1/round1_sync_position.npy' # 经纬度
position_path2 = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/round2_sync_position.npy' # 经纬度
position_path3 = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round3/round3_sync_position.npy' # 经纬度


index_map = open(index_map,'rb')
index_map = pickle.load(index_map)
video1 = cv2.VideoCapture(video_path1)
video2 = cv2.VideoCapture(video_path2)
video3 = cv2.VideoCapture(video_path3)
GNSS1 = np.load(position_path1)
GNSS2 = np.load(position_path2)
GNSS3 = np.load(position_path3)
video_width = video1.get(cv2.CAP_PROP_FRAME_WIDTH)
video_height = video1.get(cv2.CAP_PROP_FRAME_HEIGHT)

for mode in [2,1,3]:

    # mode = 2
    sampleId = 0
    resize_ratio = 0.7

    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    width = int(video_width*resize_ratio)*3
    height = int(video_height*resize_ratio)
    out = cv2.VideoWriter('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/'+'%d.avi'%mode,fourcc,30.0,(width, height))
    delay = 1
    while True:
        i = mode - 1
        frame_count = np.shape(index_map[i])[0]
        if mode == 1:
            sampleId1 = sampleId
            sampleId2 = index_map[i][sampleId][1]
            sampleId3 = index_map[i][sampleId][2]
        if mode == 2:
            sampleId1 = index_map[i][sampleId][0]
            sampleId2 = sampleId
            sampleId3 = index_map[i][sampleId][2]
        if mode == 3:
            sampleId1 = index_map[i][sampleId][0]
            sampleId2 = index_map[i][sampleId][1]
            sampleId3 = sampleId

        video1.set(cv2.CAP_PROP_POS_FRAMES, sampleId1)
        ret, frame = video1.read()
        
        frame1 = cv2.resize(frame,(int(video_width*resize_ratio),int(video_height*resize_ratio)))

        video2.set(cv2.CAP_PROP_POS_FRAMES, sampleId2)
        ret, frame = video2.read()
        frame2 = cv2.resize(frame,(int(video_width*resize_ratio),int(video_height*resize_ratio)))
        
        video3.set(cv2.CAP_PROP_POS_FRAMES, sampleId3)
        ret, frame = video3.read()
        frame3 = cv2.resize(frame,(int(video_width*resize_ratio),int(video_height*resize_ratio)))

        cv2.putText(frame1, str(sampleId1),(50,50),1,2,(0,0,255),2)
        cv2.putText(frame2, str(sampleId2),(50,50),1,2,(0,0,255),2)
        cv2.putText(frame3, str(sampleId3),(50,50),1,2,(0,0,255),2)

        if mode == 1:

            cv2.rectangle(frame1, (0,0),(int(video_width*resize_ratio),int(video_height*resize_ratio)),(0,0,255),3)
            x1 = GNSS1[sampleId1][0]
            y1 = GNSS1[sampleId1][1]

            x2 = GNSS2[sampleId2][0]
            y2 = GNSS2[sampleId2][1]

            yaw = 180 * abs(GNSS1[sampleId1][2] - GNSS2[sampleId2][2])/np.pi
            yaw = 360-yaw if yaw > 180 else yaw
            dis = np.sqrt(np.power(x1-x2,2)+np.power(y1-y2,2))
            cv2.putText(frame2, '%.3fm'%dis,(50,90),1,2,(0,0,255),2)
            cv2.putText(frame2, '%.3f degree'%yaw,(50, 130),1,2,(0,0,255),2)

            x2 = GNSS3[sampleId3][0]
            y2 = GNSS3[sampleId3][1]

            yaw = 180 * abs(GNSS1[sampleId1][2] - GNSS3[sampleId3][2])/np.pi
            yaw = 360-yaw if yaw > 180 else yaw
            dis = np.sqrt(np.power(x1-x2,2)+np.power(y1-y2,2))
            cv2.putText(frame3, '%.3fm'%dis,(50,90),1,2,(0,0,255),2)
            cv2.putText(frame3, '%.3f degree'%yaw,(50, 130),1,2,(0,0,255),2)
        if mode == 2:
            cv2.rectangle(frame2, (0,0),(int(video_width*resize_ratio),int(video_height*resize_ratio)),(0,0,255),3)
            x1 = GNSS2[sampleId2][0]
            y1 = GNSS2[sampleId2][1]

            x2 = GNSS1[sampleId1][0]
            y2 = GNSS1[sampleId1][1]

            yaw = 180 * abs(GNSS1[sampleId1][2] - GNSS2[sampleId2][2])/np.pi
            yaw = 360-yaw if yaw > 180 else yaw
            dis = np.sqrt(np.power(x1-x2,2)+np.power(y1-y2,2))
            cv2.putText(frame1, '%.3fm'%dis,(50,90),1,2,(0,0,255),2)
            cv2.putText(frame1, '%.3f degree'%yaw,(50, 130),1,2,(0,0,255),2)

            x2 = GNSS3[sampleId3][0]
            y2 = GNSS3[sampleId3][1]

            yaw = 180 * abs(GNSS3[sampleId3][2] - GNSS2[sampleId2][2])/np.pi
            yaw = 360-yaw if yaw > 180 else yaw
            dis = np.sqrt(np.power(x1-x2,2)+np.power(y1-y2,2))
            cv2.putText(frame3, '%.3fm'%dis,(50,90),1,2,(0,0,255),2)
            cv2.putText(frame3, '%.3f degree'%yaw,(50, 130),1,2,(0,0,255),2)
        if mode == 3:
            cv2.rectangle(frame3, (0,0),(int(video_width*resize_ratio),int(video_height*resize_ratio)),(0,0,255),3)
            x1 = GNSS3[sampleId3][0]
            y1 = GNSS3[sampleId3][1]

            x2 = GNSS1[sampleId1][0]
            y2 = GNSS1[sampleId1][1]

            yaw = 180 * abs(GNSS1[sampleId1][2] - GNSS3[sampleId3][2])/np.pi
            yaw = 360-yaw if yaw > 180 else yaw
            dis = np.sqrt(np.power(x1-x2,2)+np.power(y1-y2,2))
            cv2.putText(frame1, '%.3fm'%dis,(50,90),1,2,(0,0,255),2)
            cv2.putText(frame1, '%.3f degree'%yaw,(50, 130),1,2,(0,0,255),2)

            x2 = GNSS2[sampleId2][0]
            y2 = GNSS2[sampleId2][1]

            yaw = 180 * abs(GNSS3[sampleId3][2] - GNSS2[sampleId2][2])/np.pi
            yaw = 360-yaw if yaw > 180 else yaw
            dis = np.sqrt(np.power(x1-x2,2)+np.power(y1-y2,2))
            cv2.putText(frame2, '%.3fm'%dis,(50,90),1,2,(0,0,255),2)
            cv2.putText(frame2, '%.3f degree'%yaw,(50, 130),1,2,(0,0,255),2)

        frame = np.concatenate((frame1, frame2, frame3),axis=1)
        out.write(frame)

        # cv2.imshow('img', frame)
        # key = cv2.waitKey(delay)
        # if key == 32:
        #     delay = ~delay

        sampleId += 1
        if sampleId % 100 == 0:
            print(sampleId)
        if sampleId == frame_count:
            out.release()
            break


exit(0)


